import onnx
import numpy as np

# 指定你的ONNX模型文件路径
onnx_model_path = "/home/daybeha/Documents/github/ImageFeaturesAndMatching/models/SuperGlue.onnx"

# 加载ONNX模型
onnx_model = onnx.load(onnx_model_path)

# 获取模型的输入信息
input_info = onnx_model.graph.input[0]

# 获取输入的名称和形状
input_name = input_info.name
input_shape = [dim.dim_value for dim in input_info.type.tensor_type.shape.dim]

# 打印输入信息
print(f"Input name: {input_name}")
print(f"Input shape: {input_shape}")

# 如果你的输入是一个动态形状（比如包含-1），你可以使用以下方法获取详细形状信息
input_shape_dim_info = [(dim.dim_param if dim.HasField("dim_param") else dim.dim_value) for dim in input_info.type.tensor_type.shape.dim]
print(f"Detailed input shape information: {input_shape_dim_info}")


# 获取模型的所有输入信息
print("Inputs:")
for input_info in onnx_model.graph.input:
    input_name = input_info.name
    input_shape = [dim.dim_value for dim in input_info.type.tensor_type.shape.dim]
    print(f"Input name: {input_name}")
    print(f"Input shape: {input_shape}")

# 获取模型的所有输出信息
print("\nOutputs:")
for output_info in onnx_model.graph.output:
    output_name = output_info.name
    output_shape = [dim.dim_value for dim in output_info.type.tensor_type.shape.dim]
    print(f"Output name: {output_name}")
    print(f"Output shape: {output_shape}")